#include <stdio.h>
#include <string.h>
#include <arpa/inet.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <sys/time.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <unistd.h>
#include <fcntl.h>
#include <pthread.h>
#include "tls/server.h"
#include "tls/common.h"
#include "crypto/errors.h"
#include "crypto/log.h"
#include "crypto/tool.h"
#include "crypto/parse/json/ParseJsonImp.h"
#include "crypto/json.h"

#include "crypto/base64.h"
#include "crypto/des.h"
#include "crypto/rsa.h"

#define SLEEP_TIME_MS 200
#define SELECT_TIME_OUT 200

struct ConncetThreadParam
{
    int clientFd;
    struct sockaddr_in *clientAddr;
};

// 聊天数据 消息队列
struct MessageQueue
{
    int to;     // 默认群聊，只有一个聊天室 -1
    Json *data; // 转发 // "消息内容{type:"chat data", content:"", to , name, id}"
    struct MessageQueue *next;
    struct MessageQueue *pre;
};

typedef struct MessageQueue MQ;

// 头插尾入
static MQ *mqEnter = NULL;
static MQ *mqExit = NULL;
static User *clientList = NULL;
// 读写锁 增加并发
static pthread_rwlock_t clientListMutex = PTHREAD_RWLOCK_INITIALIZER;

// 由于MQ出队与入队 1:1 所以无需读写锁机制
static pthread_mutex_t messageQueueMutex = PTHREAD_MUTEX_INITIALIZER;

static int maxFd = 0;

static User *UserQuery(int fd)
{
    // 加读锁
    pthread_rwlock_rdlock(&clientListMutex);
    if (clientList == NULL)
    {
        return NULL;
    }
    User *now = clientList;
    while (now != NULL)
    {
        if (now->fd == fd)
        {
            pthread_rwlock_unlock(&clientListMutex);
            return now;
        }
    }
    pthread_rwlock_unlock(&clientListMutex);
    return NULL;
}

static MQ *MQinit(Json *json, int to)
{
    MQ *mq = malloc(sizeof(MQ));
    mq->data = json; //{type : "chat data", from : fd, name : "nickName", content : "chat data"} mq->pre = NULL;
    mq->next = NULL;
    mq->to = to;
}

static void setFDToFDS(fd_set *readFds, fd_set *exceptFds)
{
    FD_ZERO(readFds);
    FD_ZERO(exceptFds);
    pthread_rwlock_rdlock(&clientListMutex);
    User *now = clientList;
    while (now != 0)
    {
        FD_SET(now->fd, readFds);
        FD_SET(now->fd, exceptFds);
        now = now->next;
    }
    pthread_rwlock_unlock(&clientListMutex);
}

// 入队
static void MQPut(MQ *msg)
{
    pthread_mutex_lock(&messageQueueMutex);
    if (mqEnter == NULL)
    {
        mqEnter = msg;
        mqExit = msg;
    }
    else if (mqEnter == mqExit)
    {
        msg->next = mqEnter;
        mqEnter->pre = msg;
    }
    pthread_mutex_unlock(&messageQueueMutex);
}

// 出队
static MQ *MQGet()
{
    pthread_mutex_lock(&messageQueueMutex);
    MQ *tmp;
    if (mqEnter == NULL)
    {
        tmp = NULL;
    }
    else
    {
        if (mqEnter == mqExit)
        {
            tmp = mqEnter;
            mqEnter = (mqExit = NULL);
        }
        else
        {
            tmp = mqEnter;
            mqEnter = mqEnter->next;
            mqEnter->pre = NULL;
        }
    }
    pthread_mutex_unlock(&messageQueueMutex);
    return tmp;
}

static void *TcpDealClientConnect(void *args)
{
    char clientRandom[8], serverRandom[8], premasterKey[8];
    char buffer[2048], buffer2[2048];
    char *strPtr;
    int len;
    struct ConncetThreadParam *pt = args;

    /*******************************************************************************
     * 1. Client hello {type:"client hello",clientRandom:"base64 encode clientRandom data!"}
     ******************************************************************************/
    len = read(pt->clientFd, buffer, 2047);
    buffer[len] = 0;
    Json *json = ParseJson(buffer);
    SAFE_array *originalData = (SAFE_array *)malloc(sizeof(SAFE_array) + 100);
    SAFE_array *decodeData = (SAFE_array *)malloc(sizeof(SAFE_array) + 100);
    JsonValue *jsonValue;
    originalData->cap = 100;
    decodeData->cap = 100;
    decodeData->size = 0;
    JsonItem *jsonItem = JsonGet(json, "clientRandom");
    assert(jsonItem);
    strPtr = (char *)jsonItem->value->value;
    printf("Receive base64 decode clientRandom: %s\n", strPtr);
    len = strlen(strPtr);
    // printf("clientRandom: %s\n", (char *)strPtr);
    memcpy(decodeData->data, strPtr, len);
    decodeData->size = len;
    decodeBase64(decodeData, &originalData);
    memcpy(clientRandom, originalData->data, 8);

    sprintf(buffer, "[fd %d]Connnecting step 1: Client hello finish. \n", pt->clientFd);
    INFO(buffer);
    printf("%s", buffer);
    /*******************************************************************************
     * 2. Server hello {type:"server hello",serverRandom:"base64 encode serverRandom data!"}
     ******************************************************************************/
    JsonRemove(json, "clientRandom");
    generateOriginalCryptoKey(serverRandom);
    JsonPut(json, JsonStringItemInit("type", "server hello"));

    jsonValue = JsonValueInit(JSON_VALUE_TYPE_STRING);
    jsonValue->value = malloc(sizeof(20));
    memcpy(originalData->data, serverRandom, 8);
    originalData->size = 9;
    ((char *)originalData->data)[8] = 0;
    encodeBase64(originalData, &decodeData);
    memcpy(jsonValue->value, decodeData->data, decodeData->size);
    ((char *)jsonValue->value)[12] = 0;
    printf("Send base64 decode serverRandom: %s\n", (char *)(jsonValue->value));
    jsonItem = JsonItemInit("serverRandom", jsonValue);
    JsonPut(json, jsonItem);

    JsonString(json, buffer);
    len = strlen(buffer);
    write(pt->clientFd, buffer, len);

    sprintf(buffer, "[fd %d]Connnecting step 2: Server hello finish. \n", pt->clientFd);
    INFO(buffer);
    printf("%s", buffer);
    // sleep 1 s 避免多次write被一次读出
    sleep(1);
    /*******************************************************************************
     * 3. Server key exchange (RSA) {type: "server key exchange", e: 257, n: 32位大整数}
     ******************************************************************************/
    JsonRemove(json, "serverRandom");
    JsonPut(json, JsonStringItemInit("type", "server key exchange"));

    RSASuit *suit = RSAgetRSASuit(NULL);
    JsonPut(json, JsonNumberInit("e", suit->e));
    JsonPut(json, JsonNumberInit("n", suit->n));

    JsonString(json, buffer);
    len = strlen(buffer);
    write(pt->clientFd, buffer, len);

    printf(buffer, "[fd %d]Connnecting step 3: Server key exchange finish. \n", pt->clientFd);
    INFO(buffer);
    printf("%s", buffer);
    // sleep 1 s 避免多次write被一次读出
    sleep(1);
    /****************************************************************
     * 4. Server hello done {type: "server hello done"}
     *****************************************************************/
    JsonRemove(json, "e");
    JsonRemove(json, "n");
    JsonPut(json, JsonStringItemInit("type", "server hello done"));
    JsonString(json, buffer);
    len = strlen(buffer);
    write(pt->clientFd, buffer, len);

    sprintf(buffer, "[fd %d]Connnecting step 4: Server hello done finish. \n", pt->clientFd);
    INFO(buffer);
    printf("%s", buffer);
    /****************************************************************
     * 5. client key exchange (RSA) {type: "client key exchange", preMasterSecret: "rsa encode key"}
     * preMasterSecret  base64(48位) -> rsa(36位) -> 9位preMasrerSecret 前四位
     *****************************************************************/
    len = read(pt->clientFd, buffer, 2047);
    buffer[len] = 0;
    json = JsonParse(buffer);
    jsonItem = JsonGet(json, "preMasterSecret");
    originalData->size = 48;
    strcpy(originalData->data, jsonItem->value->value);
    // printf("preMasterKey: %s\n", (char *)jsonItem->value->value);
    decodeBase64(originalData, &decodeData);
    SAFE_array *rsaDecodeData = RSADecode(suit->n, suit->d, decodeData);
    memcpy(premasterKey, rsaDecodeData->data, 8);
    // printf("preMasterKey: %llu\n", *(uint64 *)premasterKey);
    free(rsaDecodeData);
    sprintf(buffer, "[System] The clientRandom is %llu. "
                    "The serverRandom is %llu. The premaster key is %llu. ",
            *(uint64 *)clientRandom, *(uint64 *)serverRandom, *(uint64 *)premasterKey);
    INFO(buffer);

    sprintf(buffer, "[fd %d]Connnecting step 5: Client key exchange. \n", pt->clientFd);
    INFO(buffer);
    printf("%s", buffer);
    /****************************************************************
     *  6. Server ChangeCipherSpec. No need TCP communication.
     *****************************************************************/
    // generate deskey
    User *user = (User *)malloc(sizeof(User));
    generateFinalDesKey(*((uint64 *)clientRandom), *((uint64 *)serverRandom), *((uint64 *)premasterKey), user->desKey);
    user->clientAddr = pt->clientAddr;
    user->fd = pt->clientFd;
    strcpy(user->nickName, "Anonymous");
    user->next = NULL;
    inet_ntop(AF_INET, &(pt->clientAddr->sin_addr), user->ip, sizeof(user->ip));
    sprintf(buffer, "[System] The deskey key is %llu", *(uint64 *)user->desKey);
    INFO(buffer);
    printf("%s\n", buffer);
    sprintf(buffer, "[fd %d]Connnecting step 6: Client ChangeCipherSpec. \n", pt->clientFd);
    INFO(buffer);
    printf("%s", buffer);
    /****************************************************************
     *  7. Receive client finish. Use Des + MD5 deal data. {type:"finished"}
     *****************************************************************/
    len = read(pt->clientFd, buffer, 2048);
    if (md5check(buffer, len))
    {
        printf("MD5 check successfully!\n");
        Json *rec = TCPReceive(buffer, len - 16, user->desKey);
        JsonItem *item = JsonGet(json, "type");
        assert(item);
        sprintf(buffer2, "client [%s]: Int connection step 7, Client TLS connection preparation is %s.", user->ip, (char *)item->value->value);
        INFO(buffer2);
        free(rec);
    }
    else
    {
        sprintf(buffer2, "client [%s]: Int connection step 7, Receive data's length is %d. MD5 check failed!", user->ip, len);
        ERROR(buffer2);
        JsonFree(json);
        free(originalData);
        free(decodeData);
        free(user);
        return NULL;
    }

    sprintf(buffer, "[fd %d]Connnecting step 7: Send client finish!\n", pt->clientFd);
    INFO(buffer);
    printf("%s", buffer);
    /****************************************************************
     *  8. Return Finish. Use Des + MD5 deal data. {type:"finished"}
     *****************************************************************/
    JsonRemove(json, "preMasterSecret");
    JsonPut(json, JsonStringItemInit("type", "finished"));
    JsonString(json, buffer);
    len = strlen(buffer);
    len = TCPSend(buffer, len, user->desKey, buffer2);
    write(user->fd, buffer2, len);

    // operation userlist, insert into header
    pthread_rwlock_wrlock(&clientListMutex);
    if (clientList == NULL)
    {
        clientList = user;
    }
    else
    {
        user->next = clientList;
        clientList = user;
    }
    maxFd = maxFd > user->fd ? maxFd : user->fd;
    printf("[INFO] MAX FD is %d.\n", maxFd);
    pthread_rwlock_unlock(&clientListMutex);
    sprintf(buffer, "[fd %d]Connnecting step 8: Receive server Finish.\n", pt->clientFd);
    INFO(buffer);
    printf("%s", buffer);
    sprintf(buffer, "[fd %d]Connnect established\n", pt->clientFd);
    INFO(buffer);
    printf("%s", buffer);
    free(pt);
    JsonFree(json);
    free(originalData);
    free(decodeData);
    return NULL;
}

// read thread.
static void *TcpServerReadThead()
{
    int retNum;
    fd_set readFds;
    fd_set exceptFds;
    int fd, ret, len;
    char buffer[2048];
    struct timeval selectTimeout;
    int count = 0;
    selectTimeout.tv_sec = 0;

    while (1)
    {
        while (clientList == NULL)
        {
            usleep(500 * 1000); // 睡眠模式 500ms judge一次
            if (count <= 0)
            {
                printf("Read thread heart beat. Maxfd is %d. Sleep state!\n", maxFd);
                count = 20; // 10s heart beat
            }
            else
            {
                count--;
            }
        }

        setFDToFDS(&readFds, &exceptFds);

        selectTimeout.tv_usec = 200 * 1000; // Select timeout is 200ms. Every time use select, the param selectTimeout will be given value 0.
        ret = select(maxFd + 1, &readFds, NULL, &exceptFds, &selectTimeout);

        if (count <= 0)
        {
            printf("Read thread heart beat. Maxfd is %d.\n", maxFd);
            count = 50; // 10s heart beat
        }
        else
        {
            count--;
        }

        if (ret == 0)
        {
            continue;
        }
        else if (ret < 0)
        {
            WARNING("In function TcpServerReadThead, the call for select return error");
        }
        else
        {
            printf("Select ret is %d\n", ret);
            pthread_rwlock_wrlock(&clientListMutex);
            User *now = clientList;
            User *pre = NULL;
            while (now != NULL)
            {
                if (FD_ISSET(now->fd, &readFds))
                {
                    len = read(now->fd, buffer, 2048);
                    if (len == 0) // 连接的读一半关闭，即接收到FIN段，读操作将返回0（对方调用close，返回为0表示对等方关闭；）
                    {
                        close(now->fd);
                        sprintf(buffer, "The remote connection %s whose fd is %d has closed", now->ip, now->fd);
                        INFO(buffer);
                        printf("%s\n", buffer);
                        if (pre == NULL)
                        {
                            free(now->clientAddr);
                            free(now);
                            clientList = now->next;
                            now = clientList;
                        }
                        else
                        {
                            pre->next = now->next;
                            free(now->clientAddr);
                            free(now);
                            now = pre;
                        }
                        if (now == NULL)
                        {
                            break;
                        }
                    }
                    else if (!md5check(buffer, len)) // md5校验
                    {
                        sprintf(buffer, "client [%s]: receive md5check error!", now->ip);
                        WARNING(buffer);
                        printf("[Debug] %s\n", buffer);
                    }
                    else
                    {
                        Json *json = TCPReceive(buffer, len - 16, now->desKey);
                        JsonItem *jsonItem = JsonGet(json, "type");
                        if (jsonItem == NULL || jsonItem->value->type != JSON_VALUE_TYPE_STRING)
                        {
                            sprintf(buffer, "client [%s]: receive data formation error! No `type` member!", now->ip);
                            WARNING(buffer);
                            printf("[Debug] %s\n", buffer);
                            free(json);
                        }
                        else
                        {
                            char *strPtr = jsonItem->value->value;
                            if (!strcmp(strPtr, "change name")) //  "chat data" {type:"change name", content:"new name"}
                            {
                                printf("[Debug] Change name command!\n");
                                jsonItem = JsonGet(json, "content");
                                if (jsonItem == NULL || jsonItem->value->type != JSON_VALUE_TYPE_STRING)
                                {
                                    sprintf(buffer, "client [%s]: Change Name but the data is error formation!", now->ip);
                                    WARNING(buffer);
                                    printf("[Debug] %s\n", buffer);
                                    free(json);
                                }
                                else
                                {
                                    strPtr = jsonItem->value->value;
                                    if (strlen(strPtr) > 19)
                                    {
                                        // cut to make the length of the new name smaller than 20.
                                        strPtr[19] = '\0';
                                    }
                                    strcpy(now->nickName, strPtr);
                                    printf("[Debug] Receive client %d data to change name to %s.", now->fd, strPtr);
                                }
                            }
                            else // "chat data" {type:"chat data", content:"chat message", to:-1}
                            {
                                printf("[Debug] chat data\n");
                                jsonItem = JsonGet(json, "content");
                                if (jsonItem == NULL || jsonItem->value->type != JSON_VALUE_TYPE_STRING)
                                {
                                    sprintf(buffer, "client [%s]: Change Name but the data is error formation!", now->ip);
                                    WARNING(buffer);
                                    free(json);
                                }
                                else
                                {
                                    jsonItem = JsonGet(json, "to");
                                    if (jsonItem->value->type != JSON_VALUE_TYPE_LONG)
                                    {
                                        sprintf(buffer, "client [%s]: Chat data does not have a member named 'to' which is a number!", now->ip);
                                        free(json);
                                        continue;
                                    }
                                    JsonPut(json, JsonStringItemInit("name", now->nickName));
                                    JsonPut(json, JsonNumberInit("id", now->fd));
                                    MQPut(MQinit(json, *(int64 *)jsonItem->value->value));
                                }
                            }
                        }
                    }
                }
                else if (FD_ISSET(now->fd, &exceptFds)) // socket 异常需要删除User
                {
                    close(now->fd); // 异常关闭fd
                    if (pre == NULL)
                    {
                        free(now);
                        clientList = now->next;
                        now = clientList;
                    }
                    else
                    {
                        pre->next = now->next;
                        free(now);
                        now = pre;
                    }

                    if (now == NULL)
                    {
                        break;
                    }
                }
                pre = now;
                now = now->next;
            }
            pthread_rwlock_unlock(&clientListMutex);
            usleep(200 * 1000); // sleep 200ms
        }
    }
}

static void *TcpServerWriteThead()
{
    MQ *msg;
    User *user;
    int len;
    char jsonStringBuffer[2048];
    byte cryptoDataBuffer[2048];
    char logBuffer[2048];
    int count = 0;
    while (1)
    {
        if ((msg = MQGet()) != NULL)
        {
            JsonString(msg->data, jsonStringBuffer); // json数据
            if (msg->to <= 0)                        // 小于0 群聊 目前只有-1所有人群聊
            {
                pthread_rwlock_rdlock(&clientListMutex);
                user = clientList;
                while (user != NULL)
                {
                    len = TCPSend(jsonStringBuffer, strlen(jsonStringBuffer), user->desKey, cryptoDataBuffer); // 每个用户的私钥不一样
                    write(user->fd, cryptoDataBuffer, len);
                    user = user->next;
                }
                pthread_rwlock_unlock(&clientListMutex);
            }
            else if (msg->to > 0) // 大于0 fd私聊
            {
                user = UserQuery(msg->to);
                if (user != NULL)
                {
                    len = TCPSend(jsonStringBuffer, strlen(jsonStringBuffer), user->desKey, cryptoDataBuffer);
                    write(user->fd, cryptoDataBuffer, len);
                }
            }
            else // 等于0 系统指令
            {
                // sprintf(logBuffer, "[System Instruction][Json] %s", jsonStringBuffer);
                // logBuffer[2047] = '\0';
                // logBuffer[2046] = '.';
                // logBuffer[2045] = '.';
                // logBuffer[2044] = '.';
                // logBuffer[2043] = '.';
                // logBuffer[2042] = '.';
                // logBuffer[2041] = '.';
                // INFO(logBuffer);
                INFO(jsonStringBuffer);
            }
            JsonFree(msg->data); // 释放Json
            free(msg);           // 释放MQ的元素
        }
        else
        {
            usleep(200 * 1000); // sleep 200ms
        }
        if (count == 0)
        {
            printf("Write thread heart beat\n");
            count = 50; // 10 heart beat
        }
        else
        {
            count--;
        }
    }
}

int TcpServer(char *ip, uint16 port, int32 BACKLOG)
{
    int fd = 0, on = 1;
    struct sockaddr_in listenAddr = {0};  // 本地地址信息结构图，下面有具体的属性赋值
    fd = socket(AF_INET, SOCK_STREAM, 0); // 建立socket
    // 默认是阻塞的
    // int flags = fcntl(fd, F_GETFL, 0);
    // fcntl(fd, F_SETFL, flags | O_SYNC);
    assert(fd > 0);
    listenAddr.sin_family = AF_INET;   /*该属性表示接收本机或其他机器传输*/
    listenAddr.sin_port = htons(port); /*端口号*/
    if (ip == NULL)
    {
        listenAddr.sin_addr.s_addr = htonl(INADDR_ANY);
    }
    else
    {
        uint32 listenAddrIp;
        inet_pton(AF_INET, ip, &listenAddrIp);
        listenAddr.sin_addr.s_addr = listenAddrIp;
    }

    assert(setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on)) >= 0);
    assert(bind(fd, (void *)&listenAddr, sizeof(listenAddr)) >= 0);
    assert(listen(fd, BACKLOG) >= 0);

    return fd;
}

int TcpServerStart(int listenFd)
{
    pthread_t th;
    pthread_create(&th, NULL, TcpServerReadThead, NULL);
    pthread_create(&th, NULL, TcpServerWriteThead, NULL);

    socklen_t clientAddrLen;
    INFO("TLS server start successfully!");
    while (true)
    {
        struct sockaddr_in *clientAddr = (struct sockaddr_in *)malloc(sizeof(struct sockaddr_in));
        clientAddrLen = sizeof(struct sockaddr_in);
        int clientFd = accept(listenFd, (struct sockaddr *)clientAddr, &clientAddrLen);
        printf("Client Connect: %s:%d\n", inet_ntoa(clientAddr->sin_addr), clientAddr->sin_port);
        fflush(stdout);
        // char buffer[20];
        // inet_ntop(AF_INET, &(clientAddr->sin_addr), buffer, sizeof(buffer));
        // printf("ip: %s\n", buffer);
        if (clientFd <= 0)
        {
            free(clientAddr);
            perror("accept");
            ERROR("Client");
        }
        else
        {
            struct ConncetThreadParam *threadParam = (struct ConncetThreadParam *)malloc(sizeof(struct ConncetThreadParam));
            threadParam->clientFd = clientFd;
            threadParam->clientAddr = clientAddr;
            pthread_create(&th, NULL, TcpDealClientConnect, threadParam);
            // TcpDealClientConnect(threadParam);
        }
    }
}